/* june 12
 * functions to integrate ODEs for likelihood.py
 * 
 * TODO memory managment;
 */

//for python
#include <Python.h>
#include <numpy/arrayobject.h>

//for gsl
#include <stdlib.h>
#include <stdio.h>
#include <math.h>

//for ODEs;  maybe switch to vode
#include <gsl/gsl_errno.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_odeiv2.h>
#include <gsl/gsl_blas.h>


/* Macros to extract array attributes.
 * http://code.google.com/p/pyamg/source/browse/trunk/Examples/SWIG/numpy.i?r=801
 */
#define is_array(a)            ((a) && PyArray_Check((PyArrayObject *)a))
#define array_type(a)          (int)(PyArray_TYPE(a))
#define array_numdims(a)       (((PyArrayObject *)a)->nd)
#define array_dimensions(a)    (((PyArrayObject *)a)->dimensions)
#define array_size(a,i)        (((PyArrayObject *)a)->dimensions[i])
#define array_data(a)          (((PyArrayObject *)a)->data)
#define array_is_contiguous(a) (PyArray_ISCONTIGUOUS(a))
#define array_is_native(a)     (PyArray_ISNOTSWAPPED(a))

//~ cProgrammingWithPythonArrays.pdf
#define IND1(a, i) *((double *)(a->data + i*a->strides[0]))
#define IND2(a, i, j) \
*((double *)(a->data + i*a->strides[0] + j*a->strides[1]))



typedef struct {
	double t0;
	double deltat;
	//birth and migration matrices; mXm
	PyArrayObject* b0;
	PyArrayObject* b1;
	PyArrayObject* mig0;
	PyArrayObject* mig1;
	//pop size vector; m
	PyArrayObject* p0;
	PyArrayObject* p1;
	// which matrix elements are nonzero?
	PyArrayObject* nonZeroElements; //numpy array; in python is a list of 2-tuple;
	int lengthNonZeroElements; 
	// number states
	int m; 
} AdotParams;

/////////////////////////
// interpolators
float linearInterpNumpyMatrix(double t, int k, int l, AdotParams *params, PyArrayObject* y0, PyArrayObject* y1){
	return IND2(y0, k,l) + ( IND2(y1, k,l) - IND2(y0, k,l)) * (t - params->t0) / (params->deltat); 
	}
float linearInterpBirths(double t, int k, int l, AdotParams *params){
	return linearInterpNumpyMatrix(t, k, l,params, params->b0, params->b1);
}
float linearInterpMigs(double t, int k, int l, AdotParams *params){
	return linearInterpNumpyMatrix(t, k, l,params, params->mig0, params->mig1);
}
float linearInterpPopSize(double t, int k, AdotParams *params){
	return IND1(params->p0, k) + ( IND1(params->p1, k) - IND1(params->p0, k)) * (t - params->t0) / (params->deltat); 
	}


// derivatives
int Adot (double t, const double A[], double dA[],
           void *params_in)
{
	AdotParams *params = (AdotParams *)params_in;
	double fractionYnotA[params->m];
	int i,k,l; 
	float popSize[params->m];
	
	
	//working:
	//~ printf("%d\n", params->m);
	//~ printf("%d\n", params->lengthNonZeroElements);
	//~ printf("%f\n", params->t0);
	//~ printf("%f\n", params->deltat);
	//~ printf("%f", IND1(params->p0,4));
	
	for (i = 0; i < params->m; i++){
		popSize[i] = linearInterpPopSize(t, i, params);
		dA[i] = 0.;
	}
	for (i = 0; i < params->m; i++){
		fractionYnotA[i] = (popSize[i] - A[i]) / popSize[i];
		if (fractionYnotA[i] < 0){
			fractionYnotA[i] = 0.;
		}
	}
	
	for (i = 0; i < params->lengthNonZeroElements; i++){
		k = IND2(params->nonZeroElements, i, 0);
		l = IND2(params->nonZeroElements, i, 1);
		if (k!=l){
			dA[k] = dA[k] + A[l] * linearInterpMigs(t,k,l,params) / popSize[l] \
				  - A[k] * linearInterpMigs(t, l, k, params) / popSize[k] \
				  + (A[l] / popSize[l]) * fractionYnotA[k] * linearInterpBirths(t,k,l,params) \
				  - (A[k] / popSize[k]) * linearInterpBirths(t,l, k,params);
		} else{
			dA[k] =dA[k] -  (A[k] / popSize[k]) * (A[k] / popSize[k]) * linearInterpBirths(t,k,l,params) ; 
		}
	}
	return GSL_SUCCESS;
}


// can perhaps make this faster by solving each of the Qi vectors separately
// and using  a spline for A
//~ #define Ak(AQS, k) AQS[k]
#define Qkl(mat, k, l, m) mat[m + k*m + l]
int AQSdot(double t, const double AQS[], double dAQS[],
           void *params_in)
{
	AdotParams *params = (AdotParams *)params_in;
	double fractionYnotA[params->m];
	
	int i,k,l; 
	float popSize[params->m];
	float A[params->m];
	
	//debug
	if (params->deltat <=0){
		printf("params->deltat %f", params->deltat);
	}
	
	for (i = 0; i < params->m; i++){
		popSize[i] = linearInterpPopSize(t, i, params);
		A[i] = AQS[i];
	}
	for (i = 0; i < params->m; i++){
		//debug
		if (popSize[i]<=0){
			printf("popsize k %f\n", popSize[i]);
			printf("A k %f\n", A[i]);
		}
		
		fractionYnotA[i] = (popSize[i] - A[i]) / popSize[i];
		if (fractionYnotA[i] < 0){
			fractionYnotA[i] = 0.;
		}
		
		//~ printf("fraction Y not A %f\n", fractionYnotA[i]);
	}
	for (i =0; i< (params->m + (params->m) * (params->m) + 1); i++){
		dAQS[i] = 0;
	}
	
	
	
	for (i = 0; i < params->lengthNonZeroElements; i++){
		k = IND2(params->nonZeroElements, i, 0);
		l = IND2(params->nonZeroElements, i, 1);
		
		
		
		dAQS[k] = dAQS[k] + A[l] * linearInterpMigs(t,k,l,params) / popSize[l] \
			  - A[k] * linearInterpMigs(t, l, k, params) / popSize[k] \
			  + (A[l] / popSize[l]) * fractionYnotA[k] * linearInterpBirths(t,k,l,params) \
			  - (A[k] / popSize[k]) * fractionYnotA[l] * linearInterpBirths(t,l, k,params);
		
	}
	
	
	
	//Qdot
	double lambdaS;
	double AlOverPl, rowsum;
	double R[params->m][params->m];//transition matrix
	//~ double **R;
	//~ R = (double **) malloc( (params->m) * sizeof(double *));
	lambdaS = 0.;
	for (k = 0; k < params->m; k++){
		for (l = 0; l < params->m; l++){
			R[k][l] = 0.;
		}
	}
	//~ printf("R 3 6 %f\n", R[3][6]);
	for (i = 0; i < params->lengthNonZeroElements; i++){
		k = IND2(params->nonZeroElements, i, 0);
		l = IND2(params->nonZeroElements, i, 1);
		if (k!=l){
			lambdaS = lambdaS + (A[k]/popSize[k]) * (A[l]/popSize[l]) * linearInterpBirths(t, k, l, params);
			R[k][l] = R[k][l] + linearInterpMigs(t,l,k,params)/popSize[k] + linearInterpBirths(t,l,k,params)*fractionYnotA[l] / popSize[k];
		}
		else if ( popSize[k] > 0 && popSize[l] > 0){
			AlOverPl = ((A[l]-1)/(popSize[l]-1));
			if (AlOverPl <= 0){
				AlOverPl = ((A[l])/(popSize[l]));
			}
			else if (AlOverPl > 1){
				AlOverPl = 1.;
			}
			lambdaS = lambdaS + (A[k]/popSize[k]) * AlOverPl * linearInterpBirths(t, k, l, params);
		}
	}
	for (k = 0; k< params->m; k++){
		rowsum = 0.;
		for (l = 0; l < params->m; l++){
			rowsum = rowsum + R[k][l];
		}
		R[k][k] = -rowsum;
		//~ printf("R kk %f\n", R[k][k]);
	}
	
	//xor code the dot product from scratch (more concise)
	double pi0[params->m];
	double sum_pi0, dpik;
	
	for (i = 0; i < params->m; i++){
		sum_pi0 = 0.;
		//normalize pi1 prior to calculating dpi
		for (k = 0; k < params->m; k++){
			// get the kth elmeent of vector i
			pi0[k] = Qkl(AQS,i,k,params->m);
			if (pi0[k] < 0) {
				//~ printf("t%f k%d pik%f\n", t, k, pi0[k]);
				pi0[k] = 0;
			}
			sum_pi0 = sum_pi0 + pi0[k];
		}
		//normalize pi0
		for (k = 0; k < params->m; k++){
			pi0[k] = pi0[k] / sum_pi0;
		}
		for (k = 0; k < params->m; k++){
			dpik = 0.;
			for (l = 0; l < params->m; l++){
				dpik = dpik + pi0[l]*R[l][k];
			}
			Qkl(dAQS,i,k,params->m) = dpik;
		}
	}
	
	
	
	//Sdot
	//~ dx[-1] = -S * lambdaS
	if (AQS[params->m + (params->m) * (params->m)] < 0){
		dAQS[params->m + (params->m) * (params->m)] = - AQS[params->m + (params->m) * (params->m)];
	}
	else{
		dAQS[params->m + (params->m) * (params->m)] = -lambdaS * AQS[params->m + (params->m) * (params->m)];
	}
	
	
	//~ free(R);
	//~ gsl_matrix_free (Rmat);
	//~ gsl_vector_free(pk);
	//~ gsl_vector_free(pk1);
	return GSL_SUCCESS;
}


////////////////////////////////////////////////////////////////////////

static PyObject *likhelper_solveA(PyObject *self, PyObject *args ){
	//unpack args: initial conditions, time limits, spline values
	double t0, deltat;
	int m, lengthNonZeroElements;
	PyArrayObject *A0, *b0, *b1, *mig0, *mig1, *p0, *p1, *nonZeroElements;
	if (!PyArg_ParseTuple(args, "ddiiOOOOOOOO", &t0, &deltat, &m, &lengthNonZeroElements, &A0, &b0, &b1, &mig0, &mig1, &p0, &p1, &nonZeroElements)) {
		PyErr_SetString(PyExc_ValueError, "wrong args.") ; 
		return NULL;  
	}
	
	//~ printf("p0 4 %f\n", IND1(p0,4));
	//~ AdotParams params = {t0, deltat, &b0, &b1, &mig0, &mig1, &p0, &p1, &nonZeroElements, lengthNonZeroElements, m};
	AdotParams params = {t0, deltat, b0, b1, mig0, mig1, p0, p1, nonZeroElements, lengthNonZeroElements, m};
	//~ printf("params.p0 4 %f\n", IND1( params.p0,4));
	
	//solve the ode's; 
	// deriv, jacobian, dimension, parameters
	gsl_odeiv2_system sys = {Adot, NULL, params.m, &params};
	
	//gsl_odeiv2_step_rk8pd
	gsl_odeiv2_driver * d = gsl_odeiv2_driver_alloc_y_new (&sys, gsl_odeiv2_step_msadams, 1e-6, 1e-6, 0.0);
	int i, status;
	double t = t0;
	double ti = t0 + params.deltat;
	double A[A0->dimensions[0]];
	for (i = 0; i < A0->dimensions[0]; i++){
		A[i] = IND1(A0, i);
	}
	status = gsl_odeiv2_driver_apply (d, &t, ti, A);
	if (status != GSL_SUCCESS) {
		  printf ("error, return value=%d\n", status);
	}
	
	
	gsl_odeiv2_driver_free (d);
	
	
	// return array 
	int dimsO[1] = {params.m};
	PyArrayObject *o = (PyArrayObject *) PyArray_FromDims(1, dimsO, PyArray_DOUBLE) ;
	//TODO can't get the return value of this to work...
	//~ PyArrayObject *o = (PyArrayObject *) PyArray_SimpleNewFromData(1, dimsO, NPY_DOUBLE, &A) ; 
	if (o == NULL) {
		printf("creating array failed\n");
		return NULL; /* PyArray_FromDims raises an exception */
	}
	for (i = 0; i < params.m; i++){
		IND1(o, i) = A[i];
	}
	
	//~ for (i = 0; i < params.m; i++){
		//~ printf("%f,%f\n", A[i], IND1(o,i));
		//~ printf("%f\n", IND1(o,i));
	//~ }
	
	//~ return (PyObject *)o;
	return PyArray_Return(o);
	//~ return Py_BuildValue("");
}




static PyObject *likhelper_solveAQS(PyObject *self, PyObject *args ){
	//unpack args: initial conditions, time limits, spline values
	double t0, deltat;
	int m, lengthNonZeroElements;
	PyArrayObject *AQS0, *b0, *b1, *mig0, *mig1, *p0, *p1, *nonZeroElements;
	if (!PyArg_ParseTuple(args, "ddiiOOOOOOOO", &t0, &deltat, &m, &lengthNonZeroElements, &AQS0, &b0, &b1, &mig0, &mig1, &p0, &p1, &nonZeroElements)) {
		PyErr_SetString(PyExc_ValueError, "wrong args.") ; 
		return NULL;  
	}
	
	AdotParams params = {t0, deltat, b0, b1, mig0, mig1, p0, p1, nonZeroElements, lengthNonZeroElements, m};
	
	//solve the ode's; 
	// deriv, jacobian, dimension, parameters
	gsl_odeiv2_system sys = {AQSdot, NULL, params.m + params.m * params.m + 1, &params};
	
	//gsl_odeiv2_step_rk8pd
	gsl_odeiv2_driver * d = gsl_odeiv2_driver_alloc_y_new (&sys, gsl_odeiv2_step_msadams, 1e-8, 1e-8, 1e-8);
	int i, status;
	double t = t0;
	double ti = t0 + params.deltat;
	double AQS1[AQS0->dimensions[0]];
	for (i = 0; i < AQS0->dimensions[0]; i++){
		AQS1[i] = IND1(AQS0, i);
	}
	status = gsl_odeiv2_driver_apply (d, &t, ti, AQS1);
	if (status != GSL_SUCCESS) {
		  printf ("error, return value=%d\n", status);
		  return NULL;
	}
	gsl_odeiv2_driver_free (d);
	
	// return array 
	int dimsO[1] = {AQS0->dimensions[0]};
	PyArrayObject *o = (PyArrayObject *) PyArray_FromDims(1, dimsO, PyArray_DOUBLE) ;
	//~ PyArrayObject *o = (PyArrayObject *) PyArray_SimpleNewFromData(1, dimsO, PyArray_DOUBLE, &AQS1) ;
	if (o == NULL) {
		printf("creating array failed\n");
		return NULL; /* PyArray_FromDims raises an exception */
	}
	for (i = 0; i < AQS0->dimensions[0]; i++){
		IND1(o, i) = AQS1[i];
	}
	
	return PyArray_Return(o);
}




//////////////////////////
static PyMethodDef likhelper_methods[] = {
  {"solveA",  (PyCFunction)likhelper_solveA, METH_VARARGS, NULL},
  {"solveAQS",  (PyCFunction)likhelper_solveAQS, METH_VARARGS, NULL},
  { NULL, NULL, 0, NULL }
};



PyMODINIT_FUNC initlikhelper(void)
{
	import_array(); // required by numpy
	Py_InitModule("likhelper", likhelper_methods);
}
